import paramiko
import hashlib
import time
import os


class Transfer:
   
    def __init__(self,ssh_args,logger=None):
        self.logger = logger
        self.ssh = paramiko.SSHClient()
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.ssh.connect(*ssh_args)
        self.sftp = self.ssh.open_sftp()
    
    def __enter__(self):
        return self  # 返回自身或其他对象

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is not None:
            print(f"Exception occurred: {exc_type}, {exc_value}")
        else:
            self.close()
    
    def makedir(self,remote_path):
        dirname = os.path.dirname(remote_path)
        try:
            self.sftp.listdir(dirname)
        except IOError as e:
            self.sftp.mkdir(dirname)
        return
    
    def compare(self,local_file,remote_file):
        local_hash = hashlib.md5(open(local_file, 'rb').read()).hexdigest()
        with self.sftp.file(remote_file, 'rb') as f:
            remote_hash = hashlib.md5(f.read()).hexdigest()
        return local_hash == remote_hash

    def put_with_progress(self,local_path,remote_path):
        name = os.path.basename(local_path)
        def print_progress(transferred, total):
            progress = round(transferred / total * 10) * 10
            if self.logger is None:
                print(f"{name}: Progress {progress:.0f}%")
            else:
                self.logger.info(f"{name}: Progress {progress:.0f}%") # , end='\r'
        temp_path = f'{remote_path}.incomplete'
        try: 
            self.sftp.put(local_path,temp_path,confirm=False,callback=print_progress)
            self.sftp.rename(temp_path, remote_path)
            return remote_path
        except IOError as e:
            time.sleep(60)
            return self.put_with_progress(local_path,remote_path)
        
    def single(self,local_path,remote_dir):
        if not os.path.exists(local_path): return
        self.makedir(remote_dir)  # 远程目录创建
        file_name = os.path.basename(local_path)
        remote_path = os.path.join(remote_dir,file_name)
        try:
            file_stat = self.sftp.stat(remote_path)
            file_size = os.path.getsize(local_path)
            if file_stat.st_size == file_size:return remote_path
            else: return self.put_with_progress(local_path,remote_path)
        except FileNotFoundError as e:
            return self.put_with_progress(local_path,remote_path)
          
    def close(self):
        self.sftp.close()
        self.ssh.close()


def core(ssh_args,remote_folder,input_queue,output_queue=None,logger=None):
    with Transfer(ssh_args,logger) as transfer:
        while True:
            element = input_queue.get()
            if element is None: break
            local_path = element
            dir_name = os.path.dirname(local_path)
            yearmonth = os.path.basename(dir_name)
            remote_dir = os.path.join(remote_folder,yearmonth)
            result = transfer.single(local_path,remote_dir)
            if result and output_queue: output_queue.put(result)
    if output_queue:output_queue.put(None)   
    return       
